import pandas
from PyQt5.QtWidgets import QDialog, QPushButton, QVBoxLayout, QTreeView, QAbstractItemView
from PyQt5.QtGui import QStandardItemModel, QStandardItem
from PyQt5 import QtCore
from PyQt5.QtCore import pyqtSignal
from collections import defaultdict
import os


class ChoiceTable(QDialog):
    selected_workbook_tables = pyqtSignal(str, dict)

    def __init__(self, filelist):
        self.chosen_workbooks_tables = defaultdict(list)
        super(ChoiceTable, self).__init__()
        self.setWindowTitle("选择表格")
        self.filelist = filelist
        layout = QVBoxLayout()
        self.tree_view = QTreeView()
        self.tree_view.header().hide()
        self.model = QStandardItemModel()
        self.load_excel(self.filelist)
        self.tree_view.setModel(self.model)
        layout.addWidget(self.tree_view)
        self.tree_view.setSelectionMode(QAbstractItemView.MultiSelection)
        self.tree_view.setEditTriggers(QTreeView.NoEditTriggers)
        selection_model = self.tree_view.selectionModel()
        # 设定选择规则
        selection_model.selectionChanged.connect(self.on_selection_changed)
        self.tree_view.expandAll()
        button = QPushButton("已经选好了")
        button1 = QPushButton("重新选择文件")
        button.clicked.connect(self.choose_done)
        button1.clicked.connect(self.close_window)
        layout.addWidget(button)
        layout.addWidget(button1)
        self.setLayout(layout)
        self.setGeometry(200, 200, 800, 600)

    def choose_done(self):
        select_index = self.tree_view.selectedIndexes()
        if len(select_index) != 0:
            for i in select_index:
                self.chosen_workbooks_tables[i.parent().data()].append(i.data())
        self.close()
        self.selected_workbook_tables.emit(os.path.split(self.filelist[0])[0], dict(self.chosen_workbooks_tables))

    def on_selection_changed(self, selected, deselected):
        indexes = selected.indexes()
        for index in indexes:
            item = self.model.itemFromIndex(index)
            if item.hasChildren():
                self.tree_view.selectionModel().select(
                    index, QtCore.QItemSelectionModel.Deselect
                )

    def load_excel(self, filelist):
        root_item = self.model.invisibleRootItem()
        for file in filelist:
            workbook = pandas.ExcelFile(file)
            workbook_item = QStandardItem(os.path.split(file)[1])
            root_item.appendRow(workbook_item)
            for sheet_name in workbook.sheet_names:
                sheet_item = QStandardItem(sheet_name)
                workbook_item.appendRow(sheet_item)
            workbook.close()

    def close_window(self):
        self.close()
